// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/functions/remote/server/remote_function_service.h>
#include <pollux/common/base/exceptions.h>
#include <pollux/expression/expr.h>
#include <pollux/functions/remote/if/get_serde.h>
#include <pollux/type/hive/hive_type_parser.h>
#include <pollux/vector/vector_stream.h>

namespace kumo::pollux::functions {
    namespace {
        std::string getFunctionName(
            const std::string &prefix,
            const std::string &functionName) {
            return prefix.empty()
                       ? functionName
                       : fmt::format("{}.{}", prefix, functionName);
        }

        TypePtr deserializeType(const std::string &input) {
            // Use hive type parser/serializer.
            return type::hive::HiveTypeParser().parse(input);
        }

        RowTypePtr deserializeArgTypes(const std::vector<std::string> &argTypes) {
            const size_t argCount = argTypes.size();

            std::vector<TypePtr> argumentTypes;
            std::vector<std::string> typeNames;
            argumentTypes.reserve(argCount);
            typeNames.reserve(argCount);

            for (size_t i = 0; i < argCount; ++i) {
                argumentTypes.emplace_back(deserializeType(argTypes[i]));
                typeNames.emplace_back(fmt::format("c{}", i));
            }
            return ROW(std::move(typeNames), std::move(argumentTypes));
        }
    } // namespace

    std::vector<core::TypedExprPtr> getExpressions(
        const RowTypePtr &inputType,
        const TypePtr &returnType,
        const std::string &functionName) {
        std::vector<core::TypedExprPtr> inputs;
        for (size_t i = 0; i < inputType->size(); ++i) {
            inputs.push_back(std::make_shared<core::FieldAccessTypedExpr>(
                inputType->childAt(i), inputType->nameOf(i)));
        }

        return {
            std::make_shared<core::CallTypedExpr>(
                returnType, std::move(inputs), functionName)
        };
    }

    void RemoteFunctionServiceHandler::handleErrors(
        apache::thrift::field_ref<remote::RemoteFunctionPage &> result,
        exec::EvalErrors *evalErrors,
        const std::unique_ptr<VectorSerde> &serde) const {
        const std::int64_t numRows = melon::copy(result->rowCount().value());
        BufferPtr dataBuffer =
                AlignedBuffer::allocate<StringView>(numRows, pool_.get());

        auto flat_vector = std::make_shared<FlatVector<StringView> >(
            pool_.get(),
            VARCHAR(),
            nullptr, // null vectors
            numRows,
            std::move(dataBuffer),
            std::vector<BufferPtr>{});

        for (vector_size_t i = 0; i < numRows; ++i) {
            if (evalErrors->hasErrorAt(i)) {
                auto exceptionPtr = *evalErrors->errorAt(i);
                try {
                    std::rethrow_exception(*exceptionPtr);
                } catch (const std::exception &ex) {
                    flat_vector->set(i, ex.what());
                }
            } else {
                flat_vector->set(i, StringView());
                flat_vector->set_null(i, true);
            }
        }
        auto errorRowVector = std::make_shared<RowVector>(
            pool_.get(),
            ROW({VARCHAR()}),
            BufferPtr(),
            numRows,
            std::vector<VectorPtr>{flat_vector});
        result->errorPayload_ref() =
                rowVectorToIOBuf(errorRowVector, *pool_, serde.get());
    }

    void RemoteFunctionServiceHandler::invokeFunction(
        remote::RemoteFunctionResponse &response,
        std::unique_ptr<remote::RemoteFunctionRequest> request) {
        const auto &functionHandle = request->remoteFunctionHandle().value();
        const auto &inputs = request->inputs().value();

        // Deserialize types and data.
        auto inputType = deserializeArgTypes(functionHandle.argumentTypes().value());
        auto outputType = deserializeType(functionHandle.returnType().value());

        auto serdeFormat = melon::copy(inputs.pageFormat().value());
        auto serde = getSerde(serdeFormat);

        auto inputVector = IOBufToRowVector(
            inputs.payload().value(), inputType, *pool_, serde.get());

        // Execute the expression.
        const vector_size_t numRows = inputVector->size();
        SelectivityVector rows{numRows};

        // Expression boilerplate.
        auto queryCtx = core::QueryCtx::create();
        core::ExecCtx execCtx{pool_.get(), queryCtx.get()};
        exec::ExprSet exprSet{
            getExpressions(
                inputType,
                outputType,
                getFunctionName(functionPrefix_, functionHandle.name().value())),
            &execCtx
        };

        exec::EvalCtx evalCtx(&execCtx, &exprSet, inputVector.get());
        if (!melon::copy(request->throwOnError().value())) {
            *evalCtx.mutableThrowOnError() = false;
        }

        std::vector<VectorPtr> expressionResult;
        exprSet.eval(rows, evalCtx, expressionResult);

        // Create output vector.
        auto outputRowVector = std::make_shared<RowVector>(
            pool_.get(), ROW({outputType}), BufferPtr(), numRows, expressionResult);

        auto result = response.result_ref();
        result->rowCount_ref() = outputRowVector->size();
        result->pageFormat_ref() = serdeFormat;
        result->payload_ref() =
                rowVectorToIOBuf(outputRowVector, rows.end(), *pool_, serde.get());

        auto evalErrors = evalCtx.errors();
        if (evalErrors != nullptr && evalErrors->hasError()) {
            handleErrors(result, evalErrors, serde);
        }
    }
} // namespace kumo::pollux::functions
